"""
  此文件对Ngsim数据集原始数据做过滤
"""

# !处理路径导入问题（添加绝对路径）！！！
import sys
import os
CODE_INTERNAL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) # 生成Code文件夹内部对应的绝对路径
sys.path.append(CODE_INTERNAL_PATH)

# 导入内部包
from utils.read_data import read_original_data
from utils.write_data import write_data_to_file
# 导入外部包
import numpy as np
from scipy import signal

# 文件路径常量
FILE_PATH_I80_1_from = "../../Data/Ngsim数据集/I80数据集/1. 原始数据/trajectories-0400-0415.txt" # 2052 884 18
FILE_PATH_I80_2_from = "../../Data/Ngsim数据集/I80数据集/1. 原始数据/trajectories-0500-0515.txt" # 1836 526 18
FILE_PATH_I80_3_from = "../../Data/Ngsim数据集/I80数据集/1. 原始数据/trajectories-0515-0530.txt" # 1790 1153 18
FILE_PATH_101_1_from = "../../Data/Ngsim数据集/101数据集/1. 原始数据/trajectories-0750am-0805am.txt" # 2169 437 18
FILE_PATH_101_2_from = "../../Data/Ngsim数据集/101数据集/1. 原始数据/trajectories-0805am-0820am.txt" # 2017 569 18
FILE_PATH_101_3_from = "../../Data/Ngsim数据集/101数据集/1. 原始数据/trajectories-0820am-0835am.txt" # 1915 421 18

FILE_PATH_I80_1_to = "../../Data/Ngsim数据集/I80数据集/2. 过滤数据/trajectories-0400-0415_filter.txt" # 2052 884 18
FILE_PATH_I80_2_to = "../../Data/Ngsim数据集/I80数据集/2. 过滤数据/trajectories-0500-0515_filter.txt" # 1836 526 18
FILE_PATH_I80_3_to = "../../Data/Ngsim数据集/I80数据集/2. 过滤数据/trajectories-0515-0530_filter.txt" # 1790 1153 18
FILE_PATH_101_1_to = "../../Data/Ngsim数据集/101数据集/2. 过滤数据/trajectories-0750am-0805am_filter.txt" # 2169 437 18
FILE_PATH_101_2_to = "../../Data/Ngsim数据集/101数据集/2. 过滤数据/trajectories-0805am-0820am_filter.txt" # 2017 569 18
FILE_PATH_101_3_to = "../../Data/Ngsim数据集/101数据集/2. 过滤数据/trajectories-0820am-0835am_filter.txt" # 1915 421 18

# 对加速度和纵向位置X做平滑处理
def filter_acc_and_x(origin_data):
  res_data = []

  for i in range(len(origin_data)):
    # origin_data是一个三维数组，取其中一个车辆ID组的数据进行平滑处理
    one_car_data = origin_data[i]
    one_car_data_np = np.array(one_car_data)

    # 过滤得到acc、x的平滑数据。采用3阶多项式拟合
    smoothed_x = signal.savgol_filter(one_car_data_np[:, 4], 21, 3).tolist()
    smoothed_acc = signal.savgol_filter(one_car_data_np[:, 12], 21, 3).tolist()

    # 检测 并 处理异常加速度
    handle_invalid_acc(smoothed_acc)

    # 将过滤好的数据更新回原始数组
    update_data(smoothed_x, one_car_data, 4)
    update_data(smoothed_acc, one_car_data, 12)

    # 基于加速度，计算新的位置和速度
    calc_new_s_and_v(one_car_data)

    res_data.append(one_car_data)
  
  return res_data

#! 处理加速度中的异常值
def handle_invalid_acc(acc_arr):
    for i in range(len(acc_arr)):
        # 对异常加速度做处理
        if is_valid_acc(acc_arr[i]):
            acc_arr[i] = round(acc_arr[i], 5)
        else:
            print("一个异常acc", i, acc_arr[i])
            # 对这个异常加速度做修改
            if i == len(acc_arr) - 1:
                acc_arr[i] = acc_arr[i - 1]
            else:
                # 基于前后加速度取均值，但是这个可能还是异常值
                temp_acc = (acc_arr[i - 1] + acc_arr[i + 1]) / 2

                while not is_valid_acc(temp_acc):
                    temp_acc = temp_acc / 2

                acc_arr[i] = round(temp_acc, 5)

# 检测是不是合法加速度
def is_valid_acc(acc):
    # 单位换算。
    acc = acc * 0.3048

    if acc < -8 or acc > 4:
        return False
    else:
        return True

# 将当前数据填充到原始数组中
def update_data(cur_arr, target_arr, idx):
    for i in range(len(cur_arr)):
        target_arr[i][idx] = cur_arr[i]

# 基于加速度，计算新的位置、速度
def calc_new_s_and_v(data_one_car, s_idx=5, v_idx=11, acc_idx=12):
    #! 0索引的s、v作为初态（s、v、acc都表示当前时刻的状态，下一时刻的s、v，由当前时刻的s、v和下一时刻的acc计算得到）
    for i in range(1, len(data_one_car)):
        cur_s = data_one_car[i - 1][s_idx]
        cur_v = data_one_car[i - 1][v_idx]
        next_acc = data_one_car[i][acc_idx]

        # s = s0 + v0t + 1/2at^2
        next_s = cur_s + cur_v * 0.1 + next_acc * 0.01 / 2
        # v = v0 + at
        next_v = cur_v + next_acc * 0.1

        # 用新数据更新原始数组
        data_one_car[i][s_idx] = round(next_s, 5)
        data_one_car[i][v_idx] = round(next_v, 5)

if __name__ == '__main__':
  # 读取数据
  origin_data = read_original_data(FILE_PATH_101_2_from)
  
  # 过滤数据（不会改变数据量大小）
  filter_data = filter_acc_and_x(origin_data)

  # 写入数据
  change_unit_idxs = [4, 5, 6, 7, 8, 9, 11, 12, 16] #! 在这一步写入时直接将单位转换好，以后不再做单位转换
#   write_data_to_file(filter_data, FILE_PATH_101_2_to, change_unit_idxs)